from .build import ARGUMENT_REGISTRY
from .compose import Compose, MapDatasetArgument
from .transform import *

ARGUMENT_REGISTRY.register('argument_randomflip', Compose(
    RandomFlip(h_prob=0.5, v_prob=0.5),
))

ARGUMENT_REGISTRY.register('argument_randomhflip', Compose(
    RandomFlip(h_prob=0.5, v_prob=0.0),
))

ARGUMENT_REGISTRY.register('argument_hflip_crop', Compose(
    RandomFlip(h_prob=0.5, v_prob=0.0),
    RandomApply(RandomCrop(height_ratio=0.5, width_ratio=0.5), p=0.5)
))

ARGUMENT_REGISTRY.register('argument_crop', Compose(
    RandomApply(RandomCrop(height_ratio=0.5, width_ratio=0.5), p=0.5)
))

ARGUMENT_REGISTRY.register('argument_flexible_crop', Compose(
    RandomApply(FlexibleRandomCrop(min_height_ratio=0.5, min_width_ratio=0.5), p=0.5)
))

ARGUMENT_REGISTRY.register('argument_ignore_LSI_KAIST_crop', MapDatasetArgument(
    {
        "LSI": None,
        "KAIST": None,
    },
    default = RandomApply(RandomCrop(height_ratio=0.5, width_ratio=0.5), p=0.5)
))

ARGUMENT_REGISTRY.register('argument_ignore_KAIST_crop', MapDatasetArgument(
    {
        "KAIST": None,
    },
    default = RandomApply(RandomCrop(height_ratio=0.5, width_ratio=0.5), p=0.5)
))

ARGUMENT_REGISTRY.register('argument_flexible_crop2', Compose(
    FlexibleRandomCrop(min_height_ratio=0.5, min_width_ratio=0.5)
))
